"""使用OpenAI API格式调用本地模型的示例"""

import openai
import os

# 配置OpenAI客户端
client = openai.OpenAI(
    base_url="http://localhost:8000/v1",  # 本地服务地址
    api_key="not_needed"  # 本地服务不需要API密钥
)

# 列出可用模型
def list_available_models():
    """列出所有可用的模型"""
    try:
        response = client.models.list()
        print("可用模型列表:")
        for model in response.data:
            print(f"- {model.id}")
        return response
    except Exception as e:
        print(f"获取模型列表失败: {str(e)}")
        return None

# 发送聊天请求
def chat_completion(prompt, model="Qwen2.5-7B-Instruct"):
    """发送聊天请求并获取响应"""
    try:
        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "你是一个有用的助手。"},
                {"role": "user", "content": prompt}
            ],
            temperature=0.7,
            max_tokens=1024
        )
        
        # 提取并返回响应内容
        answer = response.choices[0].message.content
        print(f"\n模型: {model}")
        print(f"用户: {prompt}")
        print(f"助手: {answer}")
        print(f"\n使用统计: 提示词{response.usage.prompt_tokens} tokens, 生成{response.usage.completion_tokens} tokens, 总计{response.usage.total_tokens} tokens")
        
        return response
    except Exception as e:
        print(f"聊天请求失败: {str(e)}")
        return None

# 流式响应示例
def stream_chat_completion(prompt, model="Qwen2.5-7B-Instruct"):
    """使用流式响应获取聊天结果"""
    try:
        print(f"\n模型: {model} (流式响应)")
        print(f"用户: {prompt}")
        print("助手: ", end="", flush=True)
        
        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "你是一个有用的助手。"},
                {"role": "user", "content": prompt}
            ],
            temperature=0.7,
            max_tokens=1024,
            stream=True
        )
        
        # 收集流式响应
        full_response = ""
        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                content = chunk.choices[0].delta.content
                full_response += content
                print(content, end="", flush=True)
        
        print()  # 换行
        return full_response
    except Exception as e:
        print(f"流式聊天请求失败: {str(e)}")
        return None

# 主函数
def main():
    """主函数，演示各种调用方式"""
    print("=== OpenAI兼容API调用示例 ===")
    
    # 列出可用模型
    list_available_models()
    
    # 示例1: 简单对话
    print("\n\n=== 示例1: 简单对话 ===")
    chat_completion("什么是人工智能？")
    
    # 示例2: 数学计算
    print("\n\n=== 示例2: 数学计算 ===")
    chat_completion("计算1234乘以5678的结果。")
    
    # 示例3: 代码生成
    print("\n\n=== 示例3: 代码生成 ===")
    chat_completion("用Python写一个简单的快速排序算法。")
    
    # 示例4: 流式响应
    print("\n\n=== 示例4: 流式响应 ===")
    stream_chat_completion("解释量子计算的基本原理。")
    
    # 示例5: 多轮对话
    print("\n\n=== 示例5: 多轮对话 ===")
    # 注意：实际多轮对话需要维护消息历史
    chat_completion("我想学习机器学习，你能给我一些建议吗？")

if __name__ == "__main__":
    main()